import argparse
import os
import inspect
import importlib
import compileall
import shutil
import tempfile
parser = argparse.ArgumentParser(description="生成工具")
parser.add_argument("--type",type=str, choices=['image','asr'], default='image',help="类别")
parser.add_argument("--out", default= 'out', help="项目保存目录")
parser.add_argument("--model", default= 'model.py', help="模型加载文件")
def  _removeVersion(source , dist):
    os.makedirs(dist,exist_ok=True)
    parentPath, _, files =next(os.walk(source))
    for file in files:
        name,_,ext =file.split('.')
        shutil.copyfile(os.path.join(parentPath,file),os.path.join(dist,name+"."+ext))
def build(source,dist):
    compileall.compile_dir(source,force=True)
    _removeVersion(os.path.join(source,'__pycache__') , dist)


#add  path model.py,return modelname
def getsyspath(modelpath:str):
    mpath=os.path.dirname(modelpath)
    if not mpath.startswith('/'):
        mpath = os.path.abspath(mpath)   
    os.sys.path.append(mpath)
    basename=os.path.basename(modelpath)
    name,ext = os.path.splitext(basename)
    if ext == '.py':
        return name

class GenerateImage(object):
    def __init__(self, args, work) -> None:
        self.args = args
        self.module =None
        self.work = work
    def setsyspathformodel(self):
        mpath=os.path.dirname(self.args.model)
        if not mpath.startswith('/'):
            mpath = os.path.abspath(mpath)   
        os.sys.path.append(mpath)
        basename=os.path.basename(self.args.model)
        name,ext = os.path.splitext(basename)
        if ext == '.py':
            self.module =importlib.import_module(name)


    def checkvalid(self):
        if os.path.splitext(self.args.model)[1] != '.py':
            raise Exception('模型文件错误')
        os.makedirs(self.args.out,exist_ok=True)
    def _api(self):
        data ='''from fastapi import FastAPI,Request,Response,File
from typing import List
import os
from config import getsettings
from model import getmodel
import json
import base64
import numpy as np
import cv2

app = FastAPI()
settings = getsettings()
@app.middleware("http")
async def add_process_time_header(request: Request, call_next):
    try:
        token = settings.token
    except:
        response = await call_next(request)
        return response
    try:
        if settings.token != "":      
            auth = request.headers.get("Authorization",None)
            if  auth is None:
                print("Authorization is None")
                return Response(status_code=401)
            metaname, getoken = auth.split()
            if getoken != settings.token:
                return Response(status_code=401)
        response = await call_next(request)
        return response
    except:
        return Response(status_code=401)

runurl= os.path.join(settings.baseurl,"run") if settings.baseurl else "/run"
@app.post(runurl)
async def predict(request:Request): 
    try:
        model = getmodel()
        body = await request.body()
        reqimg = img_decode(base64.b64decode(body))
        data = model(reqimg)
        return data
    except Exception as e:
        print(e)
        return Response(status_code=400)
def img_decode(content: bytes):
    np_arr = np.frombuffer(content, dtype=np.uint8)
    return cv2.imdecode(np_arr, cv2.IMREAD_COLOR)'''
        
        with open(os.path.join(self.work,"api.py"
                               ), "w") as f:
            f.write(data)
    def _config(self, extract:dict):

        if self.module is not None:
            
            data = '''from pydantic import BaseSettings
from functools import lru_cache

class Settings(BaseSettings):
'''

            cache = '''@lru_cache
def getsettings():
    set = Settings()
    return set
'''
            with open(os.path.join(self.work,"config.py"),"w") as f:
                f.write(data)
                for k,v in extract.items():
                    if isinstance(v , int):
                        f.write("\t"+k+" = "+str(v)+"\n")
                    else:
                        f.write("\t"+k+" = \""+str(v)+"\"\n")
                f.write(cache)
                  
    def _modelandconfig(self):
        if self.module is not None:
            data=inspect.getsource(self.module)
            settingsstr = inspect.getsource(self.module.settings)
            setmodel=self.module.settings.__dict__
            self._config({key:val   for key,val in setmodel.items() if not key.startswith("_")})
            settingsdata = '''from config import  getsettings
settings = getsettings()
'''
            data = data.replace(settingsstr, settingsdata)
            cache = '''from functools import lru_cache

@lru_cache
def getmodel():
    return model()
'''
            with open(os.path.join(self.work,"model.py"),"w") as f:
                f.write(data)
                f.write(cache)
    def run(self):
        try:
            self.checkvalid()
            self.setsyspathformodel()
            self._modelandconfig()
            self._api()
        except Exception as e:
            print(e)
if __name__ == '__main__':
    args = parser.parse_args()
    if args.type == "image":
        
        with tempfile.TemporaryDirectory() as tmpdirname:
            a=GenerateImage(args=args, work=tmpdirname)
            a.run()
            build(tmpdirname, args.out)





